﻿/*-
 * Copyright (C) 2018-2020, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>, 赵楷文 (Kaiwen
 * Zhao) <zhaokaiwen@ticpsh.com>
 */

#include "testif.h"

#include <clang/Basic/SourceManager.h>
#include <clang/Lex/HeaderMap.h>
#include <clang/Lex/HeaderMapTypes.h>
#include <clang/Lex/Lexer.h>
#include <clang/Lex/PPCallbacks.h>
#include <clang/Lex/Preprocessor.h>

#include <boost/filesystem.hpp>

#include "common/Coder.h"
#include "ifpreprocessor.h"

using namespace rocket::common;
using namespace clang;
using namespace std;
using njson = nlohmann::json;
namespace fs = boost::filesystem;

namespace rocket {
namespace testif {

static inline unsigned HashHMapKey(StringRef Str) {
  unsigned Result = 0;
  const char *S = Str.begin(), *End = Str.end();

  for (; S != End; S++) Result += std::tolower(*S) * 13;
  return Result;
}

class IncludeBackslash : public clang::PPCallbacks {
  clang::Preprocessor &pp;
  clang::HeaderMap hm;
  std::string current_dir;
  std::string previous_added_dir;
  std::string previous_mapped_include;

 public:
  IncludeBackslash(clang::Preprocessor &pp) : pp(pp), hm(nullptr, false) {}

  void FileChanged(clang::SourceLocation Loc,
                   clang::PPCallbacks::FileChangeReason Reason,
                   clang::SrcMgr::CharacteristicKind FileType,
                   clang::FileID PrevFID) override {
    if (!previous_added_dir.empty()) {
      auto &HS = pp.getHeaderSearchInfo();
      assert(HS.SystemDirIdx > 0);
      bool isAngled = (previous_mapped_include[0] == '<');
      if (isAngled) {
        // Fixme: angled includes cann't be mapped.
        auto it = HS.SearchDirs.begin() + HS.SystemDirIdx - 1;
        if (it->isNormalDir()) {
          SmallString<256> dir(it->getName());
          llvm::sys::fs::make_absolute(dir);
          llvm::sys::path::remove_dots(dir, /*remove_dot_dot=*/true);
          if (previous_added_dir == dir.c_str()) {
            HS.SearchDirs.erase(it);
            HS.SystemDirIdx--;
          }
        }
      } else if (HS.AngledDirIdx > 0) {
        auto it = HS.SearchDirs.begin() + HS.AngledDirIdx - 1;
        if (it->isNormalDir()) {
          SmallString<256> dir(it->getName());
          llvm::sys::fs::make_absolute(dir);
          llvm::sys::path::remove_dots(dir, /*remove_dot_dot=*/true);
          if (previous_added_dir == dir.c_str()) {
            HS.SearchDirs.erase(it);
            HS.AngledDirIdx--;
            HS.SystemDirIdx--;
          }
        }
      }
      previous_added_dir.clear();
    }

    const llvm::StringRef filename = pp.getSourceManager().getFilename(Loc);
    if (filename.empty() || filename == "/dev/null") return;

    llvm::SmallString<256> path(filename);
    llvm::sys::path::remove_filename(path);
    llvm::sys::fs::make_absolute(path);
    llvm::sys::path::remove_dots(path, /*remove_dot_dot=*/true);
    current_dir = path.c_str();
  }
#if LLVM_VERSION_MAJOR <= 14
  bool FileNotFound(llvm::StringRef FileName,
                    SmallVectorImpl<char> &RecoveryPath) override
#else
  bool FileNotFound(llvm::StringRef FileName) override
#endif
  {
    {
      std::string MappedName;
      for (auto it = FileName.begin(); it != FileName.end();) {
        char ch = *it++;
        if (ch == '\\') {
          if (MappedName.empty()) {
            MappedName.reserve(FileName.size());
            MappedName =
                StringRef(FileName.begin(), it - FileName.begin() - 1).str();
          }
          if (*it == '\\') it++;
          MappedName.push_back('/');
        } else if (!MappedName.empty()) {
          MappedName.push_back(ch);
        }
      }
      if (MappedName.empty()) return false;

      unsigned bucket = HashHMapKey(FileName);
      // number of buckets should be the least power of 2
      uint32_t NumBuckets = bucket | (bucket >> 1);
      NumBuckets = NumBuckets | (NumBuckets >> 2);
      NumBuckets = NumBuckets | (NumBuckets >> 4);
      NumBuckets = NumBuckets | (NumBuckets >> 8);
      NumBuckets = NumBuckets | (NumBuckets >> 16);
      NumBuckets += 1;
      auto buffer = llvm::WritableMemoryBuffer::getNewMemBuffer(
          sizeof(HMapHeader) + sizeof(HMapBucket) * NumBuckets + 1 +
          FileName.size() + 1 + MappedName.size() + 1);
      // header of map
      HMapHeader *header =
          reinterpret_cast<HMapHeader *>(buffer->getBufferStart());
      header->Magic = HMAP_HeaderMagicNumber;
      header->Version = HMAP_HeaderVersion;
      header->Reserved = 0;
      header->NumBuckets = NumBuckets;
      header->NumEntries = 1;
      // offset of string storage in the map
      header->StringsOffset =
          sizeof(HMapHeader) + sizeof(HMapBucket) * NumBuckets;
      // bucket of map
      HMapBucket *buckets = reinterpret_cast<HMapBucket *>(
          buffer->getBufferStart() + sizeof(HMapHeader));
      for (unsigned i = 0; i < NumBuckets; ++i) {
        buckets[i].Key = HMAP_EmptyBucketKey;
      }
      // initialize the bucket;
      buckets[bucket].Key = 1;
      char *data = buffer->getBufferStart() + header->StringsOffset +
                   buckets[bucket].Key;
      strncpy(data, FileName.data(), FileName.size());
      buckets[bucket].Prefix = buckets[bucket].Key + FileName.size() + 1;
      data += FileName.size() + 1;
      strcpy(data, MappedName.c_str());
      buckets[bucket].Suffix =
          buckets[bucket].Prefix + MappedName.size();  // empty Suffix

      hm.FileBuffer = std::move(buffer);
      {
        SmallString<1024> Path;
        assert(hm.lookupFilename(FileName, Path) == MappedName);
      }

      auto &HS = pp.getHeaderSearchInfo();

#if LLVM_VERSION_MAJOR <= 14
      // Fake path to recovery the header search of clang.
      RecoveryPath.append(current_dir.begin(), current_dir.end());
#endif

      // Erase redundant SearchDirs added by clang for
      // ReoveryPath.append(current_dir).
      if (!previous_added_dir.empty()) {
        assert(HS.SystemDirIdx > 0);
        bool isAngled = (previous_mapped_include[0] == '<');
        if (isAngled) {
          // Fixme: angled includes cann't be mapped.
          auto it = HS.SearchDirs.begin() + HS.SystemDirIdx - 1;
          if (it->isNormalDir()) {
            SmallString<256> dir(it->getName());
            llvm::sys::fs::make_absolute(dir);
            llvm::sys::path::remove_dots(dir, /*remove_dot_dot=*/true);
            if (previous_added_dir == dir.c_str()) {
              HS.SearchDirs.erase(it);
              HS.SystemDirIdx--;
            }
          }
        } else if (HS.AngledDirIdx > 0) {
          auto it = HS.SearchDirs.begin() + HS.AngledDirIdx - 1;
          if (it->isNormalDir()) {
            SmallString<256> dir(it->getName());
            llvm::sys::fs::make_absolute(dir);
            llvm::sys::path::remove_dots(dir, /*remove_dot_dot=*/true);
            if (previous_added_dir == dir.c_str()) {
              HS.SearchDirs.erase(it);
              HS.AngledDirIdx--;
              HS.SystemDirIdx--;
            }
          }
        }
      }
      previous_added_dir = current_dir;
      previous_mapped_include = FileName.str();

      // Add mapped name for header search.
      auto it = HS.SearchDirs.begin();
      if (!it->isHeaderMap() || it->getHeaderMap() != &hm) {
        DirectoryLookup DL(&hm, SrcMgr::C_User, false);
        HS.SearchDirs.insert(HS.SearchDirs.begin(), DL);
        HS.AngledDirIdx++;
        HS.SystemDirIdx++;
      }

      return true;
    }
  }
};

/**
 * @brief TestIfConsumer::HandleTopLevelDecl
 * 从顶层声明中找出待测函数及外部（全局）变量。
 * @param D
 * @return
 */
// bool TestIfConsumer::HandleTopLevelDecl(DeclGroupRef D) {
//  return true;
//}

// *************************************************************************
// 源代码遍历完成后的解析

struct MacroRange {
  clang::SourceLocation begin;
  clang::SourceLocation end;

  std::string expansion;
};

/// 同 clang/Frontend/Rewrite/RewriteMacros.cpp
/// isSameToken - Return true if the two specified tokens start have the same
/// content.
static bool isSameToken(Token &RawTok, Token &PPTok) {
  // If two tokens have the same kind and the same identifier info, they are
  // obviously the same.
  if (PPTok.getKind() == RawTok.getKind() &&
      PPTok.getIdentifierInfo() == RawTok.getIdentifierInfo())
    return true;

  // Otherwise, if they are different but have the same identifier info, they
  // are also considered to be the same.  This allows keywords and raw lexed
  // identifiers with the same name to be treated the same.
  if (PPTok.getIdentifierInfo() &&
      PPTok.getIdentifierInfo() == RawTok.getIdentifierInfo())
    return true;

  return false;
}

/// 见 clang/Frontend/Rewrite/RewriteMacros.cpp
/// GetNextRawTok - Return the next raw token in the stream, skipping over
/// comments if ReturnComment is false.
static const Token &GetNextRawTok(const std::vector<Token> &RawTokens,
                                  unsigned &CurTok, bool ReturnComment) {
  assert(CurTok < RawTokens.size() && "Overran eof!");

  // If the client doesn't want comments and we have one, skip it.
  if (!ReturnComment && RawTokens[CurTok].is(tok::comment)) ++CurTok;

  return RawTokens[CurTok++];
}

/// 同 clang/Frontend/Rewrite/RewriteMacros.cpp
/// /// LexRawTokensFromMainFile - Lets all the raw tokens from the main file
/// into the specified vector.
static void LexRawTokensFromMainFile(Preprocessor &PP,
                                     std::vector<Token> &RawTokens) {
  SourceManager &SM = PP.getSourceManager();

  // Create a lexer to lex all the tokens of the main file in raw mode.  Even
  // though it is in raw mode, it will not return comments.
#if LLVM_VERSION_MAJOR >= 12
  llvm::MemoryBufferRef FromFile = SM.getBufferOrFake(SM.getMainFileID());
  Lexer RawLex(SM.getMainFileID(), FromFile, SM, PP.getLangOpts());
#else
  const llvm::MemoryBuffer *FromFile = SM.getBuffer(SM.getMainFileID());
  Lexer RawLex(SM.getMainFileID(), FromFile, SM, PP.getLangOpts());
#endif

  // Switch on comment lexing because we really do want them.
  RawLex.SetCommentRetentionState(true);

  Token RawTok;
  do {
    RawLex.LexFromRawLexer(RawTok);

    // If we have an identifier with no identifier info for our raw token, look
    // up the identifier info.  This is important for equality comparison of
    // identifier tokens.
    if (RawTok.is(tok::raw_identifier)) PP.LookUpIdentifierInfo(RawTok);

    RawTokens.push_back(RawTok);
  } while (RawTok.isNot(tok::eof));
}

static bool isEofOrEom(tok::TokenKind Kind) {
  return Kind == tok::eof
#if LLVM_VERSION_MAJOR >= 17
         || Kind == tok::annot_module_begin || Kind == tok::annot_module_end ||
         Kind == tok::annot_module_include || Kind == tok::annot_repl_input_end
#endif
      ;
}

/// 参考 clang/Frontend/Rewrite/RewriteMacros.cpp
/// RewriteMacrosInInput - Implement -rewrite-macros mode.
static void parse_macros_expansion(Preprocessor &PP,
                                   std::vector<MacroRange> &macro_range_buf) {
  SourceManager &SM = PP.getSourceManager();

  std::vector<Token> RawTokens;
  LexRawTokensFromMainFile(PP, RawTokens);
  unsigned CurRawTok = 0;
  Token RawTok = GetNextRawTok(RawTokens, CurRawTok, false);

  // Get the first preprocessing token.
  PP.EnterMainSourceFile();
  Token PPTok;
  PP.Lex(PPTok);

  while (!isEofOrEom(RawTok.getKind()) || !isEofOrEom(PPTok.getKind())) {
    SourceLocation PPLoc = SM.getExpansionLoc(PPTok.getLocation());

    // If PPTok is from a different source file, ignore it.
    if (!SM.isWrittenInMainFile(PPLoc)) {
      PP.Lex(PPTok);
      continue;
    }

    // If the raw file hits a preprocessor directive, they will be extra tokens
    // in the raw file that don't exist in the preprocsesed file.  However, we
    // choose to preserve them in the output file and otherwise handle them
    // specially.
    if (RawTok.is(tok::hash) && RawTok.isAtStartOfLine()) {
      // Otherwise, if this is a #include or some other directive, just leave it
      // in the file by skipping over the line.
      RawTok = GetNextRawTok(RawTokens, CurRawTok, false);
      while (!RawTok.isAtStartOfLine() && RawTok.isNot(tok::eof))
        RawTok = GetNextRawTok(RawTokens, CurRawTok, false);
      continue;
    }

    // Okay, both tokens are from the same file.  Get their offsets from the
    // start of the file.
    unsigned PPOffs = SM.getFileOffset(PPLoc);
    unsigned RawOffs = SM.getFileOffset(RawTok.getLocation());

    // If the offsets are the same and the token kind is the same, ignore them.
    if (PPOffs == RawOffs && isSameToken(RawTok, PPTok)) {
      RawTok = GetNextRawTok(RawTokens, CurRawTok, false);
      PP.Lex(PPTok);
      continue;
    }

    // If the PP token is farther along than the raw token, something was
    // deleted.  Comment out the raw token.
    if (RawOffs <= PPOffs) {
      // Comment out a whole run of tokens instead of bracketing each one with
      // comments.  Add a leading space if RawTok didn't have one.
      auto EndPos = RawTok.getLocation();

      if (macro_range_buf.empty() || macro_range_buf.back().end != EndPos) {
        macro_range_buf.emplace_back();
        macro_range_buf.back().begin = EndPos;
      }

      do {
        auto RawTok_loc = RawTok.getLocation();
        if (RawTok_loc.isValid() && !RawTok.is(tok::comment)) {
          EndPos = RawTok_loc.getLocWithOffset(RawTok.getLength());
        }

        RawTok = GetNextRawTok(RawTokens, CurRawTok, true);
        RawOffs = SM.getFileOffset(RawTok.getLocation());

        if (RawTok.is(tok::comment)) {
          // Skip past the comment.
          RawTok = GetNextRawTok(RawTokens, CurRawTok, false);
          break;
        }

      } while (RawOffs <= PPOffs && !RawTok.isAtStartOfLine() &&
               (PPOffs != RawOffs || !isSameToken(RawTok, PPTok)));

      if (!macro_range_buf.empty()) {
        auto &macro_range = macro_range_buf.back();
        if (macro_range.begin == EndPos) {
          macro_range_buf.pop_back();
        } else {
          macro_range.end = EndPos;
        }
      }
      continue;
    }

    // Otherwise, there was a replacement an expansion.  Insert the new token
    // in the output buffer.  Insert the whole run of new tokens at once to get
    // them in the right order.
    std::string Expansion;
    unsigned int PPEndOffs =
        SM.getFileOffset(SM.getLocForEndOfFile(SM.getMainFileID()));
    while (PPOffs < RawOffs) {
      Expansion += (!Expansion.empty() && PPTok.hasLeadingSpace() ? " " : "") +
                   PP.getSpelling(PPTok);
      PP.Lex(PPTok);
      PPLoc = SM.getExpansionLoc(PPTok.getLocation());
      PPOffs = SM.getFileOffset(PPLoc);

      // FIXED:防止最后一个Token为宏且后跟随一个换行符时出现死循环的情况
      if (PPOffs + 1 == PPEndOffs || PPEndOffs == RawOffs) break;
    }

    if (!macro_range_buf.empty()) {
      macro_range_buf.back().expansion += Expansion;
    }
  }
}

/**
 * @brief TestIfConsumer::HandleTranslationUnit
 * 遍历变量和函数定义
 * @param ast
 */
void TestIfConsumer::HandleTranslationUnit(clang::ASTContext &ast) {
  CI->getSema().getDiagnostics().setIgnoreAllWarnings(true);
  CI->getSema().ActOnEndOfTranslationUnit();
  CI->getSema().getDiagnostics().Reset();
  CI->getSema().getDiagnostics().setIgnoreAllWarnings(false);

  if (setting.analyze_line_statistic) {
    auto &sm = ast.getSourceManager();
    // 文件首
    SourceLocation head_loc = sm.getLocForStartOfFile(sm.getMainFileID());
    // 文件尾
    SourceLocation tail_loc = sm.getLocForEndOfFile(sm.getMainFileID());
    codeline_info =
        getCodeLineInfoInRange(SourceRange(head_loc, tail_loc), ast);

    auto commentlist = CI->getASTContext().Comments.getCommentsInFile(
        CI->getSourceManager().getMainFileID());
    if (commentlist) {
      for (auto comment : *commentlist) {
        auto comment_range = comment.second->getSourceRange();
        auto comment_begin = comment_range.getBegin();
        auto comment_end = comment_range.getEnd();
        auto comment_pbegin = sm.getPresumedLoc(comment_begin);
        auto comment_pend = sm.getPresumedLoc(comment_end);

        for (unsigned l = comment_pbegin.getLine(); l <= comment_pend.getLine();
             ++l) {
          codeline_info.addCommentLine(l);
        }
      }
    }

    codeline_info.not_blank_lines = codeline_info.code_lines;
    codeline_info.not_blank_lines.insert(codeline_info.comment_lines.begin(),
                                         codeline_info.comment_lines.end());

    codelines_j["code"] = to_string(codeline_info.code_lines.size());
    codelines_j["comment"] = to_string(codeline_info.comment_lines.size());
    codelines_j["blank"] = to_string(sm.getPresumedLoc(tail_loc).getLine() -
                                     codeline_info.not_blank_lines.size());
    codelines_j["total"] = to_string(sm.getPresumedLoc(tail_loc).getLine());
  }

  TestTopLevelVisitor target(function_definitions, function_declarations,
                             global_declarations, CI->getASTContext(), setting,
                             (unsigned)setting.analyze_level);
  target.TraverseDecl(CI->getASTContext().getTranslationUnitDecl());

  /* 默认给予一个void*的类型信息 */
  parse_type(CI->getASTContext().VoidPtrTy);

  /* 解析强制要求的额外类型 */
  for (auto &type : target.types) {
    parse_type(QualType(type, 0));
  }

  /* 全局变量解析 */
  for (auto &global : global_declarations) {
    globals_j[getFullNDeclName(global.first)] = parse_variable(global.first);
  }

  /* 待测函数解析 */
  for (auto &func : function_definitions) {
    if (auto method_decl = dyn_cast<CXXMethodDecl>(func.first)) {
      if (method_decl->getParent()->isLocalClass()) {
        // 局部类的方法函数不作为待测函数
        continue;
      }
    }
    functions_j.push_back(parse_function_definition(func.first, func.second));
  }

  /* 宏展开的字符串 */
  macros_j = njson::object();
  if (setting.analyze_macro_expansions) {
    auto &macros_usages_j = macros_j["usages"] = njson::array();
    auto &PP = CI->getPreprocessor();
    if (1) {
      std::vector<MacroRange> macro_range_buf;
      parse_macros_expansion(PP, macro_range_buf);
      auto &SM = CI->getSourceManager();
      for (auto &macro_info : macro_range_buf) {
        njson macro_info_j = njson::object();

        FullSourceLoc begin(SM.getFileLoc(macro_info.begin), SM);
        FullSourceLoc end(SM.getFileLoc(macro_info.end), SM);

        {
          auto &macro_location = macro_info_j["location"] = njson::object();

          macro_location["begin"]["line"] = to_string(begin.getLineNumber());
          macro_location["begin"]["column"] =
              to_string(begin.getColumnNumber());
          macro_location["end"]["line"] = to_string(end.getLineNumber());
          macro_location["end"]["column"] = to_string(end.getColumnNumber());
        }

        macro_info_j["@expansion"] = macro_info.expansion;

        macros_usages_j.push_back(std::move(macro_info_j));
      }
    } else {
      auto expansion = parseMacrosExpansion(PP);
      auto &replaces = expansion.replaces;
      auto &SM = CI->getSourceManager();

      for (auto &replace : replaces) {
        njson macro_info_j = njson::object();

        FullSourceLoc loc(SM.getFileLoc(replace.first), SM);

        {
          auto &macro_location = macro_info_j["location"] = njson::object();
          macro_location["begin"]["line"] = to_string(loc.getLineNumber());
          macro_location["begin"]["column"] = to_string(loc.getColumnNumber());
          macro_location["end"]["line"] = to_string(loc.getLineNumber());
          macro_location["end"]["column"] = to_string(loc.getColumnNumber());
        }

        macro_info_j["@expansion"] = replace.second.expansion;

        macros_usages_j.push_back(std::move(macro_info_j));
      }
    }
  }

  /* 解析过程中出现的绝对地址的解析 */
  parse_fixed_addresses();

  auto &sm = CI->getASTContext().getSourceManager();
  for (auto type2files_it = type2files_map.begin(),
            type2files_end = type2files_map.end();
       type2files_it != type2files_end; ++type2files_it) {
    auto &fis = type2files_it->second;
    auto &path_j = *type2files_it->first;
    for (auto fi : fis) {
      if (fi.isValid() && sm.getFileEntryForID(fi)) {
#if LLVM_VERSION_MAJOR > 18
        path_j.push_back(sm.getFileEntryForID(fi)->tryGetRealPathName());
#else
        path_j.push_back(sm.getFileEntryForID(fi)->getName());
#endif
      }
    }
  }
}

bool TestIfConsumer::isInExcludeNamespace(const Decl *decl) {
  for (auto &en : setting.exclude_namespace) {
    if (common::isInNamespace(en, decl)) {
      return true;
    } else if (isCXX() && en == "std") {
      return common::isInSystemspace(decl);
    }
  }
  return false;
}

void TestIfConsumer::Initialize(clang::ASTContext &) {
  CI->getPreprocessor().addPPCallbacks(
      std::make_unique<IncludeBackslash>(CI->getPreprocessor()));
  return;
}

}  // namespace testif
}  // namespace rocket
